# Author: Ivan Canay | Northwestern University | January 2013
# Copyright (c) 2013 Northwestern University. All rights reserved.
#-------------------------------------------------------------------
# X-plain: 	This file has several function in R that are used in the
#          	joint project with Azeem M. Shaikh and Joe P. Romano.
#			"Randomization Test under a Weak Convergence Assumption"
#-------------------------------------------------------------------
# List of functions included in this file:
#  Functions for the Randomization Test
#   (0) permut.test   : Computes randomization test with permutations
#	(1) random.test   : Computes randomization test
#   (2) random.G      : Get groups of transformations
# We use the R compiler to improve the speed of the functions.
#-------------------------------------------------------------------
# Required Libraries:
require("plm");
require("Matrix");
require("data.table");
require("compiler")
require("MASS")
require(sandwich, quietly = TRUE)
require(lmtest, quietly = TRUE)

require("writexl")

#-------------------------------------------------------------------------------------------------------------
#------------------------------ FUNCTIONS FOR THE RANDOMIZATION TEST -----------------------------------------
#-------------------------------------------------------------------------------------------------------------
#-------------------------------------------------------------------
pre.permut.test  <- function(X,n,G_p,alpha=0.05){
#-------------------------------------------------------------------
# X-plain: Computes randomization critical value - PERMUTATION VERSION
#-------------------------------------------------------------------
# INPUTS: - X: this has the "data": n obs of Z followed by m obs of Y
#         - G: the group of all trasformations (permutations)
#         - alpha: significance level
#------------------------------------------------------------------
# RETURNS:- rule: binary decision of the test
#         - Nrule: binary decision of the non-randomized test
#         - cv: the critival value of the randomization test.
#         - pv: the p-value of the test.
#------------------------------------------------------------------
	N       <- length(X);   # Total sample size N = n + m
	m       <- N-n;
	M_p		<- dim(G_p)[2];   # # of elements in G

    ObsT <- abs(mean(X[1:n])-mean(X[(n+1):N]));  # observed Test stat
	NewX <- matrix(X[G_p],N,M_p);		    # transformed data
	# Compute New Test Stat over transformed data
	NewT <- abs(colMeans(NewX[1:n,])-colMeans(NewX[(n+1):N,]));
	NewT <- sort(NewT);        # sort the vector of Test Stat
	k    <- M_p- floor(M_p*alpha); # find index for quantile
	Mplus<- sum(NewT>NewT[k]); # M+ from LR book
	M0   <- sum(NewT==NewT[k]);# M0 from LR book
	prob <- (M_p*alpha-Mplus)/M0;# term a(x) from LR book
	# randomized Test
	rule <- (ObsT>NewT[k]) + (ObsT==NewT[k])*(runif(1)<=prob);
	# Non non-randomized Test
	Nrule <- (ObsT>NewT[k]);
	# Compute the p-value (formula (15.5) from LR book
	p.value<- sum(NewT>=ObsT)/M_p;

	# List of returns
	list(rule=rule, Nrule=Nrule, cv=NewT[k], pv=p.value)
 }
#-------------------------------------------------------------------
permut.test <- cmpfun(pre.permut.test);
#-------------------------------------------------------------------

#-------------------------------------------------------------------
pre.random.test  <- function(X,G,alpha=0.05){
#-------------------------------------------------------------------
# X-plain: Computes randomization critical value
#-------------------------------------------------------------------
# INPUTS: - X: this has the "data": q point estimators in our paper
#         - G: the group of all trasformations (from random.G)
#         - alpha: significance level
#------------------------------------------------------------------
# RETURNS: - rule: binary decision of the test
#         - Nrule: binary decision of the non-randomized test
#         - cv: the critival value of the randomization test.
#         - pv and pv2: p-values according to folmulae 15.5 and 15.7
#------------------------------------------------------------------
	n       <- length(X);   # sample size
	M		<- dim(G)[2];   # # of elements in G

    ObsT <- sqrt(n)*abs(mean(X))/sd(X); # observed Test stat
	NewX <- G*X;					    # transformed data
	# Compute New Test Stat over transformed data
	NewT <- abs(sqrt(n)*apply(NewX,2,mean)/apply(NewX,2,sd));
	NewT <- sort(NewT);        # sort the vector of Test Stat
	k    <- M- floor(M*alpha); # find index for quantile
	Mplus<- sum(NewT>NewT[k]); # M+ from LR book
	M0   <- sum(NewT==NewT[k]);# M0 from LR book
	prob <- (M*alpha-Mplus)/M0;# term a(x) from LR book
	# randomized Test
	rule <- (ObsT>NewT[k]) + (ObsT==NewT[k])*(runif(1)<=prob);
	# Non non-randomized Test
	Nrule <- (ObsT>NewT[k]);
	# Compute the p-value (formula (15.5) from LR book
	p.value<- sum(NewT>=ObsT)/M;
	# Compute the p-value (formula (15.7) from LR book
	p.value2<- (sum(NewT>=ObsT)+1)/(M+1);

	# List of returns
	list(rule=rule, Nrule=Nrule, cv=NewT[k], pv=p.value,pv2=p.value2)
 }
#-------------------------------------------------------------------
random.test <- cmpfun(pre.random.test);
#-------------------------------------------------------------------

#-------------------------------------------------------------------
pre.random.G  <- function(q,B=1000){
#-------------------------------------------------------------------
# X-plain: Computes the group transformations for the randomization test
#-------------------------------------------------------------------
# INPUTS: - q: the dimension of the data (q point estimators)
#         - B: the number of random draws if q>10.
#------------------------------------------------------------------
# RETURNS: - G: a B \times q matrix.
#------------------------------------------------------------------
  if(q<11){e   <- c(1,-1); B<-2^q;
           perm<-t(expand.grid(rep(list(e),q)));
  }
  # Draw B transformations from G otherwise
  else {perm <- matrix(2*((runif(q*B)>1/2)-1/2),q,B);
  }
  return(perm)
}
#-------------------------------------------------------------------
random.G <- cmpfun(pre.random.G);
#-------------------------------------------------------------------


#-------------------------------------------------------------------------------------------------------------
#-------------------------------------- FUNCTIONS FOR WILD BOOTSTRAP -----------------------------------------
#-------------------------------------------------------------------------------------------------------------
wildboot <- function(data, outcome, testvar, control, clustvar, fe = 0, numerrcorr = 0) {
#-------------------------------------------------------------------
# X-plain: Computes wild bootstrap
#-------------------------------------------------------------------
# INPUTS: - data: dataframe containing all variables used in estimation, appropriately named
#         - outcome: outcome variable as a string
#         - testvar: variables to test as a string
#         - control: control variables as a vector of strings 
#         - clustvar: variable to cluster on, as a string. Note that cluster number must run consecutively from 1 to Q  
#         - fe = 0 or 1. If 1 then include cluster fixed effects. 
#         - numerrcorr: tolerance for numerical error in computing bootstrap p-values. Because bootstrap samples are constructed from residuals,
#                       estimation results differ on at the magnitude of ~1e-15 even when the permutation is the identity. 
#------------------------------------------------------------------
# RETURNS: list containing: 
#          (1) estimates: studentized and unstudentized t-statistics and associated p-values
#          (2) ols_coef: remaining coefficients from OLS estimation
#------------------------------------------------------------------
  
  # data <- read.dta("Leader_Group_AER_2014_oldversion.dta") 
  # data <- data[!is.na(data$pct),] # remove data with missing outcome
  
  # outcome <- "zakaibag"; testvar <- "treated"; control <- c("semarab", "semrel"); clustvar <- "group_bg"; 
  # control <- c("leq", "las" ,control); fe <- 1, testvar <- "leqef"
  # data <- data; fe <- 0; numerrcorr <- 0; control <- c(control_1, year_fe)
  # testvar <- testvar_joint; control <- c(control_joint, year_fe); fe <- 0; numerrcorr <- 1e-12
  
  
  set.seed(123)
  
  if (fe == 0){
    factor_fe <- NULL
    factor_fe_boot <- NULL
  } else {
    factor_fe <- "as.factor(data[[clustvar]])"
    factor_fe_boot <- "as.factor(boot_data[[clustvar]])"
  }
  
  
  model <- reformulate(termlabels = c(testvar, control, factor_fe), response = outcome)
  model_restricted <- reformulate(termlabels = c(control, factor_fe), response = outcome)
  model_bootstrap <- reformulate(termlabels = c(testvar, control, factor_fe_boot), response = "outcome_star")
  
  n <- length(data[[outcome]])
  
  ols <- lm(model, data = data)
  coef(ols)
  
  t_n <- abs(sqrt(n)*coef(ols)[testvar])
  
  variances <- wildboot_variance(ols$residuals, data, testvar, control, clustvar, fe)
  Sigma_n <- solve(variances$Omega)%*%variances$V_n%*%solve(variances$Omega)
  t_n_stud <- t_n / sqrt(Sigma_n[1,1])

  ols_restricted <- lm(model_restricted, data = data)
  
  Q <- length(unique(data[[clustvar]]))
  permutations <- t(random.G(Q))
  G <- dim(permutations)[1]

  boot_stat <- rep(0, G)
  boot_stat_stud <- rep(0, G)
  
  for(g in 1:G){
    boot_residuals <- ols_restricted$residuals
    for (q in 1:Q){
      boot_residuals[data[[clustvar]] == q] <- boot_residuals[data[[clustvar]] == q]*permutations[g,q]
    }
    
    outcome_star <- ols_restricted$fitted.values + boot_residuals
    boot_data <- data.frame(outcome_star, data)
    boot_ols <- lm(model_bootstrap, data = boot_data)
    t_b <- abs(sqrt(n)*coef(boot_ols)[testvar])
    
    variances_b <- wildboot_variance(boot_ols$residuals, boot_data, testvar, control, clustvar, fe)
    Sigma_b <- solve(variances_b$Omega)%*%variances_b$V_n%*%solve(variances_b$Omega)
    t_b_stud <- t_b / sqrt(Sigma_b[1,1])
    
    boot_stat[g] <- t_b
    boot_stat_stud[g] <- t_b_stud
  }

  boot_crit_5 <- quantile(boot_stat, probs = c(0.95))
  boot_crit_10 <- quantile(boot_stat, probs = c(0.90))
  boot_pval <- sum(boot_stat >= t_n-numerrcorr)/G
  
  boot_crit_5_stud <- quantile(boot_stat_stud, probs = c(0.95))
  boot_crit_10_stud <- quantile(boot_stat_stud, probs = c(0.90))
  boot_pval_stud <- sum(boot_stat_stud >= t_n_stud-numerrcorr)/G
  
  clust_pval_n <- (1-pnorm(t_n_stud))*2
  clust_pval_t <- ((1-pt(t_n_stud, Q-1))*2)
  
  if (fe == 0){
    fe_dfc <- 0 
  } else {
    fe_dfc <- Q-1
  }
  L <- length(control) + fe_dfc + 1 + length(testvar) # include testvar and constant
  t_n_robust <- t_n/(sqrt(diag(variances$robust)[testvar]*n))
  robust_pval_t <- ((1-pt(t_n_robust, n-L))*2)
  
  estimates <- as.numeric(c(sum(coef(ols)[testvar]), t_n, t_n_stud, boot_pval, boot_pval_stud, clust_pval_n, clust_pval_t, t_n_robust, robust_pval_t, boot_crit_10, boot_crit_5, boot_crit_10_stud, boot_crit_5_stud))
  names(estimates) <- c("Coef", "t-stat", "t_stat_stud", "boot_pval", "boot_pval_stud", "clust_pval_n", "clust_pval_t", "t_n_robust", "robust_pval_t", "boot_crit_10", "boot_crit_5", "boot_crit_10_stud", "boot_crit_5_stud")
  
  return(list(estimates = estimates, ols_coef = coef(ols)))
}
#-------------------------------------------------------------------
wildboot_variance <- function(resid, data, testvar, control, clustvar, fe) {
#-------------------------------------------------------------------
# X-plain: Computes the value of Omega_{Z_tilde, n} and V_n as defined in equation (15)
#-------------------------------------------------------------------
# INPUTS: - resid: residuals from a previous OLS estimation of outcome on testvar + control
#         - data: dataframe containing all variables used in estimation, appropriately named
#         - testvar: variables to test as a string
#         - control: control variables as a vector of strings 
#         - clustvar: variable to cluster on, as a string. Note that cluster number must run consecutively from 1 to Q  
#         - fe = 0 or 1. If 1 then include cluster fixed effects. 
#------------------------------------------------------------------
# RETURNS:  list containing: 
#           - (1) Omega
#           - (2) V_n
#------------------------------------------------------------------
  
  # resid <- ols$residuals

  if (fe == 0){
    model_w <- ~ 1
  } else {
    model_w <- reformulate(termlabels = c("as.factor(data[[clustvar]])"))
  }
  model_z <- reformulate(termlabels = c(testvar, control, -1))
  
  
  n <- length(data[[testvar[1]]])
  Q <- length(unique(data[[clustvar]]))
  d_z <- length(attr(terms(model_z), "term.labels"))
  
  W_full <- model.matrix(model_w, data = data)
  Z_full <- model.matrix(model_z, data = data)
  Pi_full <- solve(t(W_full)%*%W_full)%*%t(W_full)%*%Z_full
  
  covn <- numeric()
  covnq <- 1 # covnq keeps track of the start point of each cluster in the Z_tilde matrix
  Z_tilde <- numeric()
  epsilon <- numeric() # reorder residuals to match entries in Z_tilde
  for (q in 1:Q){
    W_q <- W_full[data[[clustvar]] == q, ]
    Z_q <- Z_full[data[[clustvar]] == q, ]
    Z_tilde_q <- Z_q - W_q%*%Pi_full
    n_q <- dim(Z_tilde_q)[1]
    
    epsilon_q <- resid[data[[clustvar]] == q]
    
    Z_tilde <- rbind(Z_tilde, Z_tilde_q)
    epsilon <- c(epsilon, epsilon_q)
    
    covn <- c(covn, n_q)
    covnq <- c(covnq, covnq[length(covnq)]+n_q)
  }
  
  Omega <- t(Z_tilde)%*%Z_tilde/n
  
  V_n <- Omega*0
  for (q in 1:Q){
    Z_tilde_q <- Z_tilde[covnq[q]:(covnq[q+1]-1),]
    epsilon_q <- epsilon[covnq[q]:(covnq[q+1]-1)]
    
    V_n <- V_n + t(Z_tilde_q)%*%epsilon_q%*%t(t(Z_tilde_q)%*%epsilon_q)
  }
  V_n <- V_n/n
  
  # DFC as in BCH11 pg. 142. or as on p.16 in NBER WP 18478 by Imbens and Kolesar
  # dfc <- (S/(S-1))*(dof=="bch")+(N/(N-L))*(dof=="HC1")+((N-1)/(N-L))*(S/(S-1))*(dof=="stata");
  if (fe == 0){
    fe_dfc <- 0 
  } else {
    fe_dfc <- Q-1
  }
  L <- length(control) + fe_dfc + 1 + length(testvar) # include testvar and constant
  dfc <- ((n-1)/(n-L))*(Q/(Q-1))
  # print(dfc)
  V_n <- dfc*V_n
  
  
  # Compute robust standard error
  if (fe == 0){
    factor_fe <- NULL
  } else {
    factor_fe <- "as.factor(data[[clustvar]])"
  }
  model_full <- reformulate(termlabels = c(testvar, control, factor_fe), response = outcome)
  X_full <- model.matrix(model_full, data = data)
  robust <- solve(t(X_full)%*%X_full)%*% t(X_full)%*%diag(resid^2)%*%X_full %*%solve(t(X_full)%*%X_full)
  robust <- robust*n/(n-L)
  
  #sqrt(diag(robust))
  
  #sqrt(sum(diag((robust*n))[testvar]))
  
  output <- list(Omega = Omega, V_n = V_n, robust = robust)
  
  return(output)
}
#-------------------------------------------------------------------
wildboot_clustcov <- function(data, Z, W, clustvar, fe = 0) {
  #-------------------------------------------------------------------
  # X-plain: Computes the value of the cluster covariances term in A2.2(iii)
  #-------------------------------------------------------------------
  # INPUTS: - data: dataframe containing all variables used in estimation, appropriately named
  #         - Z: variables to test, as a vector strings
  #         - W: control variables as a vector of strings. Do not include cluster fixed effects here; use FE option. 
  #         - clustvar: variable to cluster on, as a string. Note that cluster number must run consecutively from 1 to Q  
  #         - fe = 0 or 1. If 1 then include cluster fixed effects. 
  #------------------------------------------------------------------
  # RETURNS:  list containing: 
  #           - (1) covmats: list of all cluster variance/covariance matrices
  #           - (2) variances: variance of each cell across the cluster variance/covariance matrices
  #           - (3) abdsdiff: largest absolute difference for each cell across the cluster variance/covariance matrices
  #           - (4) absmax: cluster from which the largest value of each cell is observed, value of which is used in calculation of absdiff
  #           - (5) absmin: cluster from which the smallest value of each cell is observed, value of which is used in calculation of absdiff
  #------------------------------------------------------------------
  
  # Z <- c("lgrain_pred_famdum5860", "lgrain_pred"); W <- c("lurbpop", "ltotpop", year_fe); fe <- 1
  
  if (fe == 0){
    factor_fe <- NULL
  } else {
    factor_fe <- "as.factor(data[[clustvar]])"
  }
  model_w <- reformulate(termlabels = c(W, factor_fe))
  model_z <- reformulate(termlabels = c(Z, -1))
  
  # n <- length(data[[testvar]])
  Q <- length(unique(data[[clustvar]]))
  d_z <- length(attr(terms(model_z), "term.labels"))
  
  W_full <- model.matrix(model_w, data = data)
  Z_full <- model.matrix(model_z, data = data)
  Pi_full <- solve(t(W_full)%*%W_full)%*%t(W_full)%*%Z_full
  
  covcube <- (1:(d_z*d_z*Q))*0  # store all the variance-covariance matrices in a cube; index 1 refers to cluster number
  dim(covcube) <- c(Q, d_z, d_z) 
  covmats <- list()   # list of variance-covariance matrices. 
  covnames <- numeric()
  covn <- numeric() 
  
  for (q in 1:Q){
    W_q <- W_full[data[[clustvar]] == q, ]
    Z_q <- Z_full[data[[clustvar]] == q, ]
    Z_tilde_q <- Z_q - W_q%*%Pi_full
    n_q <- dim(Z_tilde_q)[1]
    mat_q <- t(Z_tilde_q)%*%Z_tilde_q/n_q
    covmats <- append(covmats, list(mat_q))
    covnames <- c(covnames, (q))
    covcube[q,,] <- mat_q
    covn <- c(covn, n_q)
  }
  
  names(covmats) <- covnames
  
  variances <- apply(covcube, c(2,3), var)
  absdiff <- apply(covcube, c(2,3), max) - apply(covcube, c(2,3), min)
  absmax <- apply(covcube, c(2,3), which.max)
  absmin <- apply(-covcube, c(2,3), which.max)
  output <- list(covmats = covmats, variances = variances, absdiff = absdiff, absmax = absmax, absmin = absmin, covcube = covcube, covn = covn)
  
  return(output)
  
}
#-------------------------------------------------------------------
normalize <- function(data, exclude){
#-------------------------------------------------------------------
# X-plain: normalize all variables to have variance 1
#-------------------------------------------------------------------
# INPUTS:   - data: dataframe to be normalized
#           - exclude: name of variables to NOT be normalized, as a vector of strings
#------------------------------------------------------------------
# RETURNS:  NULL; Saves name.xlsx in working directory. 
#------------------------------------------------------------------  
  data_norm <- data[, !names(data) %in% exclude] 
  stddevs <- apply(data_norm, 2, sd)
  n <- dim(data_norm)[1]
  stddevs <- rep(1,n)%*%t(stddevs)
  data_norm <- data_norm / stddevs
  data_keep <- data[, names(data) %in% exclude, drop=F]
  data_out <- cbind(data_norm, data_keep)
  
  colnames(data_out) <- c(colnames(data_norm), colnames(data_keep))

  print("The resulting variances are:")
  print(apply(data_out, 2, var))
  
  return(data_out)
  
}
#-------------------------------------------------------------------
print2Excel <- function(output, name, digits = 3){
#-------------------------------------------------------------------
# X-plain: prints output of clustcov to xlsx files. 
#-------------------------------------------------------------------
# INPUTS:   - output: output from wildboot_clustcov
#           - name: name of file, including xlsx suffix, e.g. "output.xlsx
#           - digits: number of digits to round output to
#------------------------------------------------------------------
# RETURNS:  NULL; Saves name.xlsx in working directory. 
#------------------------------------------------------------------  

  # output <- col_1_ols_nofe_mat
  
  d_z = dim(output[["covmats"]][[1]])[1]
  Q <- length(output[["covmats"]])
  
  ncol <- (d_z+1)
  nrow <- (d_z+2)
  block <- (matrix('',nrow = nrow,ncol = ncol))
  
  
  # We assemble blocks each starting with size (d_z+2)*(d_z+1), before a column of row names is attached. The blocks are then joint one after another. 
  var_block <- block
  var_block[2:(d_z+1), 1:d_z] <- round(output$variances, digits)
  var_block <- cbind(c("", row.names(output[["covmats"]][[1]]), ""), var_block)
  var_block[1,1] <- "Variances"
  var_block[1,2:(d_z+1)] <- row.names(output[["covmats"]][[1]])
  
  
  absdiff_block <- block
  absdiff_block[2:(d_z+1), 1:d_z] <- round(output$absdiff, digits)
  absdiff_block <- cbind(c("", row.names(output[["covmats"]][[1]]), ""), absdiff_block)
  absdiff_block[1,1] <- "Abs Diff"
  absdiff_block[1,2:(d_z+1)] <- row.names(output[["covmats"]][[1]])
  
  absmax_block <- block
  absmax_block[2:(d_z+1), 1:d_z] <- round(output$absmax, digits)
  absmax_block <- cbind(c("", row.names(output[["covmats"]][[1]]), ""), absmax_block)
  absmax_block[1,1] <- "Abs Max"
  absmax_block[1,2:(d_z+1)] <- row.names(output[["covmats"]][[1]])
  
  absmin_block <- block
  absmin_block[2:(d_z+1), 1:d_z] <- round(output$absmin, digits)
  absmin_block <- cbind(c("", row.names(output[["covmats"]][[1]]), ""), absmin_block)
  absmin_block[1,1] <- "Abs Min"
  absmin_block[1,2:(d_z+1)] <- row.names(output[["covmats"]][[1]])
  
  absdiff_block <- rbind(absdiff_block, absmax_block, absmin_block)  
  
  m <- rbind(var_block, absdiff_block)
  
  for (q in 1:Q){
    q_block <- block
    q_block[2:(d_z+1), 1:d_z] <- round(output[["covmats"]][[q]], digits)
    q_block <- cbind(c("", row.names(output[["covmats"]][[1]]), ""), q_block)
    q_block[1,1] <- paste("Cluster", q)
    q_block[1,2:(d_z+1)] <- row.names(output[["covmats"]][[1]])
    m <- rbind(m, q_block)
  }
  
  
  m <- rbind(matrix("", nrow = 2, ncol = ncol+1 ), m)
  m[1,1] <- name
  
  
  write_xlsx(x = data.frame(m), path = name, col_names=FALSE)
}
wildboot_joint <- function(data, outcome, testvar, control, clustvar, fe = 0, numerrcorr = 0) {
  #-------------------------------------------------------------------
  # X-plain: Computes wild bootstrap
  #-------------------------------------------------------------------
  # INPUTS: - data: dataframe containing all variables used in estimation, appropriately named
  #         - outcome: outcome variable as a string
  #         - testvar: variables to test as a string
  #         - control: control variables as a vector of strings 
  #         - clustvar: variable to cluster on, as a string. Note that cluster number must run consecutively from 1 to Q  
  #         - fe = 0 or 1. If 1 then include cluster fixed effects. 
  #         - numerrcorr: tolerance for numerical error in computing bootstrap p-values. Because bootstrap samples are constructed from residuals,
  #                       estimation results differ on at the magnitude of ~1e-15 even when the permutation is the identity. 
  #------------------------------------------------------------------
  # RETURNS: list containing: 
  #          (1) estimates: studentized and unstudentized t-statistics and associated p-values
  #          (2) ols_coef: remaining coefficients from OLS estimation
  #------------------------------------------------------------------
  
  # data <- read.dta("Leader_Group_AER_2014_oldversion.dta") 
  # data <- data[!is.na(data$pct),] # remove data with missing outcome
  
  # outcome <- "ldeath_b"
  # testvar <- c("lgrain_pred_famdum5860", "lgrain_pred"); clustvar <- "prov_clust"
  # control <- c("lurbpop", "ltotpop", year_fe); fe <- 0; numerrcorr <- 0
  # testvar <- testvar_joint; control <- c(control_joint, year_fe); fe <- 0; numerrcorr <- 1e-12
  
  set.seed(123)
  
  if (fe == 0){
    factor_fe <- NULL
    factor_fe_boot <- NULL
  } else {
    factor_fe <- "as.factor(data[[clustvar]])"
    factor_fe_boot <- "as.factor(boot_data[[clustvar]])"
  }
  
  
  model <- reformulate(termlabels = c(testvar, control, factor_fe), response = outcome)
  model_restricted <- reformulate(termlabels = c(control, factor_fe), response = outcome)
  model_bootstrap <- reformulate(termlabels = c(testvar, control, factor_fe_boot), response = "outcome_star")
  
  n <- length(data[[outcome]])
  
  ols <- lm(model, data = data)
  coef(ols)
  
  t_n <- abs(sqrt(n)*sum(coef(ols)[testvar]))
  
  
  variances <- wildboot_variance(ols$residuals, data, testvar, control, clustvar, fe)
  
  
  Sigma_n <- solve(variances$Omega)%*%variances$V_n%*%solve(variances$Omega)
  hyp_r <- c(rep(1.0, length(testvar)), rep(0.0, length(control)))
  t_n_stud <- as.numeric(t_n / sqrt(t(hyp_r)%*%Sigma_n%*%hyp_r))
  
  
  ols_restricted <- lm(model_restricted, data = data)
  
  Q <- length(unique(data[[clustvar]]))
  permutations <- t(random.G(Q))
  G <- dim(permutations)[1]
  
  boot_stat <- rep(0, G)
  boot_stat_stud <- rep(0, G)
  
  for(g in 1:G){
    boot_residuals <- ols_restricted$residuals
    for (q in 1:Q){
      boot_residuals[data[[clustvar]] == q] <- boot_residuals[data[[clustvar]] == q]*permutations[g,q]
    }
    
    outcome_star <- ols_restricted$fitted.values + boot_residuals
    boot_data <- data.frame(outcome_star, data)
    boot_ols <- lm(model_bootstrap, data = boot_data)
    t_b <- abs(sqrt(n)*sum(coef(boot_ols)[testvar]))
    
    variances_b <- wildboot_variance(boot_ols$residuals, boot_data, testvar, control, clustvar, fe)
    Sigma_b <- solve(variances_b$Omega)%*%variances_b$V_n%*%solve(variances_b$Omega)
    t_b_stud <- as.numeric(t_b / sqrt( t(hyp_r)%*%Sigma_b%*%hyp_r ))
    
    boot_stat[g] <- t_b
    boot_stat_stud[g] <- t_b_stud
  }
  
  
  boot_crit_5 <- quantile(boot_stat, probs = c(0.95))
  boot_crit_10 <- quantile(boot_stat, probs = c(0.90))
  boot_pval <- sum(boot_stat >= (t_n-numerrcorr))/G
  
  boot_crit_5_stud <- quantile(boot_stat_stud, probs = c(0.95))
  boot_crit_10_stud <- quantile(boot_stat_stud, probs = c(0.90))
  boot_pval_stud <- sum(boot_stat_stud >= (t_n_stud-numerrcorr))/G
  
  # if (fe == 0){
  #   fe_dfc <- 0 
  # } else {
  #   fe_dfc <- Q-1
  # }
  # L <- length(control) + fe_dfc + 1 + length(testvar) # include testvar and constant
  
  clust_pval_n <- (1-pnorm(t_n_stud))*2
  clust_pval_t <- ((1-pt(t_n_stud, Q-1))*2)
  
  if (fe == 0){
    fe_dfc <- 0 
  } else {
    fe_dfc <- Q-1
  }
  L <- length(control) + fe_dfc + 1 + length(testvar) # include testvar and constant
  hyp_robust <- c(0, rep(1, length(testvar)), rep(0, L - 1 - length(testvar)))
  t_n_robust <- as.numeric(t_n/(sqrt(  t(hyp_robust)%*%variances$robust%*%hyp_robust*n)))
  robust_pval_t <- ((1-pt(t_n_robust, n-L))*2)
  
  estimates <- as.numeric(c(sum(coef(ols)[testvar]), t_n, t_n_stud, boot_pval, boot_pval_stud, clust_pval_n, clust_pval_t, t_n_robust, robust_pval_t, boot_crit_10, boot_crit_5, boot_crit_10_stud, boot_crit_5_stud))
  names(estimates) <- c("Coef", "t-stat", "t_stat_stud", "boot_pval", "boot_pval_stud", "clust_pval_n", "clust_pval_t", "t_n_robust", "robust_pval_t", "boot_crit_10", "boot_crit_5", "boot_crit_10_stud", "boot_crit_5_stud")
  
  return(list(desc = c("Test jointly 0:", paste(testvar, collapse=", ")), estimates = estimates, ols_coef = coef(ols)))
}


############################################################
# Old Code
############################################################
# wildboot_clustcov <- function(data, testvar, control, clustvar, fe = 0) {
#   #-------------------------------------------------------------------
#   # X-plain: Computes the value of the cluster covariances term in A2.2(iii). W is assumed to be cluster fixed effects if included, constant otherwise
#   #-------------------------------------------------------------------
#   # INPUTS: - data: dataframe containing all variables used in estimation, appropriately named
#   #         - testvar: variables to test as a string
#   #         - control: control variables as a vector of strings 
#   #         - clustvar: variable to cluster on, as a string. Note that cluster number must run consecutively from 1 to Q  
#   #         - fe = 0 or 1. If 1 then include cluster fixed effects. 
#   #------------------------------------------------------------------
#   # RETURNS:  list containing: 
#   #           - (1) covmats: list of all cluster variance/covariance matrices
#   #           - (2) variances: variance of each cell across the cluster variance/covariance matrices
#   #           - (3) abdsdiff: largest absolute difference for each cell across the cluster variance/covariance matrices
#   #           - (4) absmax: cluster from which the largest value of each cell is observed, value of which is used in calculation of absdiff
#   #           - (5) absmin: cluster from which the smallest value of each cell is observed, value of which is used in calculation of absdiff
#   #------------------------------------------------------------------
#   
#   
#   # outcome <- "zakaibag"; testvar <- "treated"; control <- c("semarab", "semrel"); clustvar <- "group_bg"; fe <- 1
#   
#   if (fe == 0){
#     model_w <- ~ 1
#   } else {
#     model_w <- reformulate(termlabels = c("as.factor(data[[clustvar]])"))
#   }
#   model_z <- reformulate(termlabels = c(testvar, control, -1))
#   
#   # n <- length(data[[testvar]])
#   Q <- length(unique(data[[clustvar]]))
#   d_z <- length(attr(terms(model_z), "term.labels"))
#   
#   W_full <- model.matrix(model_w, data = data)
#   Z_full <- model.matrix(model_z, data = data)
#   Pi_full <- solve(t(W_full)%*%W_full)%*%t(W_full)%*%Z_full
#   
#   covcube <- (1:(d_z*d_z*Q))*0  # store all the variance-covariance matrices in a cube; index 1 refers to cluster number
#   dim(covcube) <- c(Q, d_z, d_z) 
#   covmats <- list()   # list of variance-covariance matrices. 
#   covnames <- numeric()
#   covn <- numeric() 
#   
#   for (q in 1:Q){
#     W_q <- W_full[data[[clustvar]] == q, ]
#     Z_q <- Z_full[data[[clustvar]] == q, ]
#     Z_tilde_q <- Z_q - W_q%*%Pi_full
#     n_q <- dim(Z_tilde_q)[1]
#     mat_q <- t(Z_tilde_q)%*%Z_tilde_q/n_q
#     covmats <- append(covmats, list(mat_q))
#     covnames <- c(covnames, (q))
#     covcube[q,,] <- mat_q
#     covn <- c(covn, n_q)
#   }
#   
#   names(covmats) <- covnames
#   
#   variances <- apply(covcube, c(2,3), var)
#   absdiff <- apply(covcube, c(2,3), max) - apply(covcube, c(2,3), min)
#   absmax <- apply(covcube, c(2,3), which.max)
#   absmin <- apply(-covcube, c(2,3), which.max)
#   output <- list(covmats = covmats, variances = variances, absdiff = absdiff, absmax = absmax, absmin = absmin, covcube = covcube, covn = covn)
#   
#   return(output)
#   
# }
